#pragma once
#include <vector>
#include <cstring>
#include <stdexcept>
#include <iostream>

using dtype = double;

enum class data_type { LINEAR, MATRIX };

class data {
    public:
    data_type type;
    dtype* vector = nullptr;
    data(const data&) = delete;
    data& operator=(const data&) = delete;
    data(data_type type) : type(type) {}
    virtual int size() const = 0;
    // virtual bool matches(data& other) const = 0;
    void set0() { if(vector == nullptr) throw std::runtime_error("set0 on nullptr"); memset(vector, '\0', sizeof(dtype) * size()); }
    void allocate() { if(vector != nullptr) delete[] vector; vector = new dtype[size()]; }
    void free() { if(vector != nullptr) delete[] vector; vector = nullptr; }
    ~data() { if(vector != nullptr) delete[] vector; }
};

class data_linear : public data {
    public:
    int length;
    data_linear(int length) : data(data_type::LINEAR), length(length) {}
    int size() const override { return length; }
    dtype& operator[](int i) const { return vector[i]; }
    // bool matches(data& other) const override { return other.type == type && ((data_linear&)other).length == length; }
};

class data_matrix : public data {
    public:
    int width, height;
    data_matrix(int width, int height) : data(data_type::MATRIX), width(width), height(height) {}
    int size() const override { return width * height; }
    dtype& get(int i, int j) const { return vector[i * height + j]; }
    void matmul(const data_linear& vect, data_linear& dst) const;
    // bool matches(data& other) const override { return other.type == type && ((data_matrix&)other).width == width && ((data_matrix&)other).height == height; }
};

enum class grad_type { DENSE, LINCOMB, LINCOMB_MODIFIED, SCALAR };

class grad {
    public:
    grad_type type;
    int in_size, out_size;
    dtype* vector = nullptr;
    grad(int in_size, int out_size, grad_type type) : in_size(in_size), out_size(out_size), type(type) {}
    grad(const grad&) = delete;
    grad& operator=(const grad&) = delete;
    virtual dtype get(int in_var, int out_var) const = 0;
    virtual bool is_allocated() const = 0;
    void combine(const grad& outs_grad, grad& dst) const;
    virtual void _combine(const grad& outs_grad, grad& dst) const = 0;
    virtual grad* get_combine_dst(const grad& outs_grad) const = 0;
    virtual int get_vector_size() const = 0;
    void allocate() { if(vector != nullptr) delete[] vector; vector = new dtype[get_vector_size()]; }
    void free() { if(vector != nullptr) delete[] vector; vector = nullptr; }
    ~grad() { if(vector != nullptr) delete[] vector; };
};

class scalar_grad : public grad {
    public:
    scalar_grad(int size) : grad(size, size, grad_type::SCALAR) {}
    dtype get(int in_var, int out_var) const override { return in_var == out_var ? vector[in_var] : 0; }
    bool is_allocated() const { return vector != nullptr; }
    void _combine(const grad& _outs_grad, grad& _dst) const override;
    grad* get_combine_dst(const grad& outs_grad) const override;
    int get_vector_size() const override { return in_size; }
};

class dense_grad : public grad {
    public:
    dense_grad(int in_size, int out_size) : grad(in_size, out_size, grad_type::DENSE) {}
    dtype get(int in_var, int out_var) const override { return vector[out_var * in_size + in_var]; }
    bool is_allocated() const { return vector != nullptr; }
    void _combine(const grad& _outs_grad, grad& _dst) const override;
    dense_grad* get_combine_dst(const grad& outs_grad) const override;
    int get_vector_size() const override { return in_size * out_size; }
};

class lincomb_grad : public grad {
    public:
    int vars;
    lincomb_grad(int vars, int combinations) : grad(vars * combinations, combinations, grad_type::LINCOMB), vars(vars) {};
    dtype get(int in_var, int out_var) const override { return in_var / vars == out_var ? vector[in_var % vars] : 0; }
    bool is_allocated() const { return vector != nullptr; }
    void _combine(const grad& _outs_grad, grad& _dst) const override;
    grad* get_combine_dst(const grad& outs_grad) const override;
    int get_vector_size() const override { return vars; }
};

class lincomb_modified_grad : public grad {
    public:
    int vars;
    lincomb_modified_grad(int vars, int combinations) : grad(vars * combinations, combinations, grad_type::LINCOMB_MODIFIED), vars(vars) {};
    dtype get(int in_var, int out_var) const override { return in_var / vars == out_var ? vector[in_var] : 0; }
    bool is_allocated() const { return vector != nullptr; }
    void _combine(const grad& _outs_grad, grad& _dst) const override;
    grad* get_combine_dst(const grad& outs_grad) const override;
    int get_vector_size() const override { return in_size; }
};
